import tensorflow as tf
import numpy as np
import pandas as pd
try:
    from tensorflow.keras import layers
except:
    from tensorflow.python.keras import layers

try:
    from tensorflow import keras
except:
    from tensorflow.python import keras

# import keras

def get_data(batch = 32):
    dataset_path = keras.utils.get_file("auto-mpg.data",
                                        "http://archive.ics.uci.edu/ml/machine-learning-databases/auto-mpg/auto-mpg.data")
    column_names = ['MPG', 'Cylinders', 'Displacement', 'Horsepower', 'Weight',
                    'Acceleration', 'Model Year', 'Origin']
    raw_dataset = pd.read_csv(dataset_path, names=column_names,
                              na_values="?", comment='\t',
                              sep=" ", skipinitialspace=True)
    dataset = raw_dataset.copy()
    dataset.head(10) # 看前10条
    dataset.isna().sum()  # 统计空白数据
    dataset = dataset.dropna()  # 删除空白数据项
    dataset.isna().sum()  # 再次统计空白数据
    origin = dataset.pop('Origin')
    dataset['USA'] = (origin == 1) * 1.0
    dataset['Europe'] = (origin == 2) * 1.0
    dataset['Japan'] = (origin == 3) * 1.0
    dataset.tail()
    train_dataset = dataset.sample(frac=0.8, random_state=0)
    test_dataset = dataset.drop(train_dataset.index)
    # 移动 MPG 油耗效能这一列为真实标签 Y
    train_labels = train_dataset.pop('MPG')
    test_labels = test_dataset.pop('MPG')
    # 查看训练集的输入 X 的统计数据
    train_stats = train_dataset.describe()
    # train_stats.pop("MPG")
    train_stats = train_stats.transpose()
    def norm(x):
        return (x - train_stats['mean']) / train_stats['std']

    normed_train_data = norm(train_dataset)
    normed_test_data = norm(test_dataset)

    print(normed_train_data.shape, train_labels.shape)
    print(normed_test_data.shape, test_labels.shape)
    # (314, 9)(314, )  # 训练集共 314 行，输入特征长度为 9,标签用一个标量表示
    # (78, 9)(78, )  # 测试集共 78 行，输入特征长度为 9,标签用一个标量表示
    # #利用切分的训练集数据构建数据集对象：
    train_db = tf.data.Dataset.from_tensor_slices((normed_train_data.values,
                                                   train_labels.values))  # 构建 Dataset 对象
    train_db = train_db.shuffle(100).batch(batch)  # 随机打散，批量化
    return train_db


class Network(keras.Model):

    # 回归网络
    def __init__(self):
        super(Network, self).__init__()

        # 创建 3 个全连接层
        self.fc1 = layers.Dense(64, activation='relu')
        self.fc2 = layers.Dense(64, activation='relu')
        self.fc3 = layers.Dense(1)

    def call(self, inputs, training=None, mask=None):

        # 依次通过 3 个全连接层
        x = self.fc1(inputs)
        x = self.fc2(x)
        x = self.fc3(x)
        return x


def train():
    model = Network()  # 创建网络类实例
    # 通过 build 函数完成内部张量的创建，其中 4 为任意的 batch 数量，9 为输入特征长度
    model.build(input_shape=(4, 9))
    model.summary()  # 打印网络信息
    optimizer = tf.keras.optimizers.RMSprop(0.001)  # 创建优化器，指定学习率
    # 接下来实现网络训练部分。通过  Epoch  和  Step  的双层循环训练网络，共训练  200    个 epoch:
    train_db = get_data(4)
    for epoch in range(200):  # 200 个 Epoch
        for step, (x, y) in enumerate(train_db):  # 遍历一次训练集
            # 梯度记录器
            with tf.GradientTape() as tape:
                out = model(x)  # 通过网络获得输出
                loss = tf.reduce_mean(tf.losses.MSE(y, out))  # 计算 MSE
                mae_loss = tf.reduce_mean(tf.losses.MAE(y, out))  # 计算 MAE
                if step % 10 == 0:  # 打印训练误差
                    print(epoch, step, float(loss))
                # 计算梯度，并更新
            grads = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(grads, model.trainable_variables))

if __name__ == '__main__':
    train()